"""
visualize.py
可视化模型预测结果
"""

import os
import random
import torch
import numpy as np
import matplotlib.pyplot as plt
from torchvision import transforms
from torchvision.datasets import CIFAR100
from PIL import Image
from model import get_model
from utils import load_checkpoint
from config import Config
from dataset import get_cifar100_superclass_mapping, get_dataloaders

# 设置随机种子
random.seed(42)
np.random.seed(42)
torch.manual_seed(42)

def get_fine_grained_classes():
    """返回CIFAR100所有细粒度类名"""
    return [
        'apple', 'aquarium_fish', 'baby', 'bear', 'beaver', 'bed', 'bee', 'beetle', 
        'bicycle', 'bottle', 'bowl', 'boy', 'bridge', 'bus', 'butterfly', 
        'camel', 'can', 'castle', 'caterpillar', 'cattle', 'chair', 'chimpanzee', 
        'clock', 'cloud', 'cockroach', 'couch', 'crab', 'crocodile', 'cup', 
        'dinosaur', 'dolphin', 'elephant', 'flatfish', 'forest', 'fox', 'girl', 
        'hamster', 'house', 'kangaroo', 'keyboard', 'lamp', 'lawn_mower', 
        'leopard', 'lion', 'lizard', 'lobster', 'man', 'maple_tree', 'motorcycle', 
        'mountain', 'mouse', 'mushroom', 'oak_tree', 'orange', 'orchid', 'otter', 
        'palm_tree', 'pear', 'pickup_truck', 'pine_tree', 'plain', 'plate', 
        'poppy', 'porcupine', 'possum', 'rabbit', 'raccoon', 'ray', 'road', 
        'rocket', 'rose', 'sea', 'seal', 'shark', 'shrew', 'skunk', 'skyscraper', 
        'snail', 'snake', 'spider', 'squirrel', 'streetcar', 'sunflower', 
        'sweet_pepper', 'table', 'tank', 'telephone', 'television', 'tiger', 
        'tractor', 'train', 'trout', 'tulip', 'turtle', 'wardrobe', 'whale', 
        'willow_tree', 'wolf', 'woman', 'worm'
    ]

def inverse_normalize(tensor, mean, std):
    """反归一化图像以便可视化"""
    for t, m, s in zip(tensor, mean, std):
        t.mul_(s).add_(m)
    return tensor

def visualize_predictions(num_images=9):
    """可视化模型预测结果"""
    cfg = Config()
    
    # 加载模型
    model = get_model()
    load_checkpoint(model)
    model.eval()
    
    # 获取数据
    _, test_loader = get_dataloaders()
    dataset = test_loader.dataset
    
    # 获取类别名称映射
    superclass_mapping = get_cifar100_superclass_mapping()
    fine_classes = get_fine_grained_classes()
    
    # 从测试集中随机选择图像
    indices = random.sample(range(len(dataset)), num_images)
    
    # 创建可视化
    plt.figure(figsize=(15, 15))
    for i, idx in enumerate(indices):
        image, true_superclass = dataset[idx]
        
        # 反归一化图像
        inv_normalize = transforms.Normalize(
            mean=[-0.5071/0.2675, -0.4867/0.2565, -0.4408/0.2761],
            std=[1/0.2675, 1/0.2565, 1/0.2761]
        )
        image = inv_normalize(image)
        
        # 预测
        with torch.no_grad():
            output = model(image.unsqueeze(0).to(cfg.DEVICE))
            _, pred_superclass = output.max(1)
            pred_superclass = pred_superclass.item()
        
        # 获取真实和预测的细粒度类
        true_fine_classes = superclass_mapping[true_superclass]
        pred_fine_classes = superclass_mapping[pred_superclass]
        
        # 转换为PIL图像并绘制
        image = transforms.ToPILImage()(image.cpu())
        
        plt.subplot(3, 3, i+1)
        plt.imshow(image)
        
        # 标题显示预测结果
        title_color = 'green' if true_superclass == pred_superclass else 'red'
        title = f"Pred: {pred_fine_classes[0]}...\nTrue: {true_fine_classes[0]}..."
        plt.title(title, color=title_color, fontsize=10)
        plt.axis('off')
        
        # 添加文本信息
        info = (f"Superclass:\n"
                f"  Pred: {pred_superclass} ({pred_fine_classes[0]}...)\n"
                f"  True: {true_superclass} ({true_fine_classes[0]}...)\n"
                f"All fine classes:\n"
                f"  Pred: {', '.join(pred_fine_classes[:3])}...\n"
                f"  True: {', '.join(true_fine_classes[:3])}...")
        plt.text(0.5, -0.3, info, transform=plt.gca().transAxes, 
                fontsize=8, ha='center', va='top')
    
    plt.tight_layout()
    plt.savefig(os.path.join(cfg.LOG_DIR, 'predictions_visualization.png'))
    plt.show()

if __name__ == "__main__":
    visualize_predictions(num_images=9)